package com.jt.dify.service.impl;

import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.jt.dify.config.LLMResponse;
import com.jt.dify.config.LLMResponseEnum;
import com.jt.dify.model.ChatRecord;
import com.jt.dify.model.ChatRecordInfo;
import com.jt.dify.model.table.ChatRecordTableDef;
import com.jt.dify.service.AiService;
import lombok.SneakyThrows;
import org.noear.solon.ai.chat.ChatModel;
import org.noear.solon.ai.chat.ChatResponse;
import org.noear.solon.ai.chat.message.ChatMessage;
import org.noear.solon.annotation.Component;
import org.noear.solon.annotation.Inject;
import org.noear.solon.net.websocket.WebSocket;
import org.noear.solon.rx.SimpleSubscriber;
import org.reactivestreams.Publisher;

import java.util.Optional;


/**
 * @author by HaiTao.Wang on 2025/2/10.
 */
@Component
public class AiServiceImpl implements AiService {


    @Inject
    ChatModel chatModel;

    @Inject("${solon.ai.dify.user}")
    String difyUser;

    @SneakyThrows
    @Override
    public void chat(WebSocket socket, String msg) {

        JSONObject parseObj = JSONUtil.parseObj(msg);
        String chatId = parseObj.getStr("id");
        String ask = parseObj.getStr("ask");

        ChatRecord one = new ChatRecord().where(ChatRecordTableDef.CHAT_RECORD.ID.eq(chatId)).one();
        if (one == null) {
            one = new ChatRecord().setContent("新对话").saveOpt().orElseThrow(RuntimeException::new).oneById();
        }

        //处理聊天记录信息
        if (one.getContent().contains("新对话")) {
            one.setContent(ask);
        }

        //构建prompt
        ChatMessage chatMessage = ChatMessage.ofUser(ask);
        chatMessage.addMetadata("user", difyUser);
        if (StrUtil.isNotEmpty(one.getConversationId())) {
            chatMessage.addMetadata("conversationId", one.getConversationId());
        }

        //chatRecordInfo信息
        ChatRecordInfo chatRecordInfo = new ChatRecordInfo().setRecordId(one.getId().toString()).setAsk(ask);

        //请求dify
        Publisher<ChatResponse> publisher = chatModel.prompt(chatMessage).stream();

        //拼接stream answer
        StringBuilder sb = new StringBuilder();

        ChatRecord finalOne = one;
        publisher.subscribe(new SimpleSubscriber<ChatResponse>()
                .doOnNext(chatResponse -> {

                    if (chatResponse.isFinished()) {

                        String fileNames = Optional.ofNullable(chatResponse.getChoices().get(0).getMessage().getMetadata().get("file")).orElse("").toString();

                        if (StrUtil.isNotEmpty(fileNames)) {
                            //报文
                            String content = JSONUtil.toJsonStr(LLMResponse.of()
                                    .setType(LLMResponseEnum.FILE)
                                    .setMsg(JSONUtil.createObj()
                                            .set("id", finalOne.getId().toString())
                                            .set("content", fileNames)
                                            .set("type", "chat").toString()));
                            socket.send(content);
                        }

                        //处理conversationId
                        if (StrUtil.isEmpty(finalOne.getConversationId())) {
                            finalOne.setConversationId(chatResponse.getMessage().getMetadata().get("conversationId").toString());
                        }

                        //更新ChatRecordInfo信息
                        chatRecordInfo.setFileName(fileNames);

                    } else {

                        sb.append(chatResponse.getMessage().getContent());
                        //报文
                        String content = JSONUtil.toJsonStr(LLMResponse.of()
                                .setType(LLMResponseEnum.ANSWER)
                                .setMsg(JSONUtil.createObj()
                                        .set("id", finalOne.getId().toString())
                                        .set("content", chatResponse.getMessage().getContent())
                                        .set("type", "chat").toString()));
                        socket.send(content);
                    }

                }).doOnComplete(() -> {

                    //报文
                    String content = JSONUtil.toJsonStr(LLMResponse.of()
                            .setType(LLMResponseEnum.FINISH)
                            .setMsg(JSONUtil.createObj()
                                    .set("id", finalOne.getId().toString())
                                    .toString()));

                    socket.send(content);

                    //更新ChatRecord
                    finalOne.updateById();

                    //新增ChatRecordInfo
                    chatRecordInfo.setAnswer(sb.toString()).save();
                }).doOnError(err -> {
                    System.err.println(err.getMessage());
                })
        );

    }

}
